use crate::{
    transport::{McpMessage, McpTransport}, McpClientCapabilities, McpError, McpErrorCode, McpImplementation, McpNotification, McpRequest, McpRequestId, McpResult, McpServerCapabilities, McpTool, McpToolCall, McpToolCallResult
};
use std::sync::Arc;
use futures::StreamExt;
use log::debug;
use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender};
use tokio::sync::{RwLock, Mutex};

pub struct McpClient {
    transport: Arc<dyn McpTransport>,
    server_capabilites: Arc<RwLock<Option<McpServerCapabilities>>>,
    request_counter: Arc<RwLock<i64>>,
    response_receiver: Arc<Mutex<UnboundedReceiver<McpMessage>>>,
    _response_sender: UnboundedSender<McpMessage>,
}

impl McpClient {
    pub fn new(transport: Arc<dyn McpTransport>) -> Self {
        // 创建一个多输出，单输入的通道
        let (tx, rx) = tokio::sync::mpsc::unbounded_channel();

        // 创建客户端
        let clinet = Self {
            transport: transport.clone(),
            server_capabilites: Arc::new(RwLock::new(None)),
            request_counter: Arc::new(RwLock::new(0)),
            response_receiver: Arc::new(Mutex::new(rx)), // 这里保存了这个 mpsc 的接收端，即由 tx 发送的消息会被 clinet 接收！
            _response_sender: tx.clone(),
        };

        // 创建一个后台任务，从 transport 中获得一个接受端（transport 是单输出，多输入的）
        // 接着不断从这个端口尝试获取 McpMessage
        // 一但接收到 McpMessage，用 mpsc 的输出端发送，就会被 clinet 的 response_receiver 收到！
        let transport_clone = transport.clone();
        let tx_clone = tx.clone();
        tokio::spawn(async move {
            let mut stream = transport_clone.receive();
            while let Some(result) = stream.next().await {
                match result {
                    Ok(message) => {
                        if tx_clone.send(message).is_err() {
                            break;
                        }
                    }
                    Err(_) => break,
                }
            }
        });

        clinet
    }

    pub async fn request(
        &self,
        method: &str,
        params: Option<serde_json::Value>,
    ) -> McpResult<serde_json::Value> {
        debug!("Request '{}' with '{:?}'", method, params);
        let mut counter = self.request_counter.write().await;
        let id = McpRequestId::Number(*counter);
        *counter += 1;

        // 整理为 McpRequest，并且调用 transport 发送出去！
        let request = McpRequest::new(method, params, id.clone());
        self.transport.send(McpMessage::Request(request)).await?;

        // 等待服务器的响应
        // 首先获取 response_receiver 的锁
        let mut receiver = self.response_receiver.lock().await;

        // 等待接受 message
        while let Some(message) = receiver.recv().await {
            // 只处理 Response 类型的消息
            debug!("Receiver message '{:?}'", message);
            if let McpMessage::Response(response) = message {
                // Id 必须匹配
                if response.id == id {
                    // 判断是否接受到错误，返回 Error
                    if let Some(error) = response.error {
                        let code: McpErrorCode = error.code.into();
                        return Err(McpError::protocol(code, &error.message))
                    }
                    // 正常接受到回复信息，取出其中的 Value 返回
                    return response.result.ok_or_else(|| McpError::protocol(
                        McpErrorCode::InternalError, 
                        "Response missing result",
                    ));
                }
            }
        }
        
        Err(McpError::protocol(
            McpErrorCode::InternalError,
            "Connection closed while waiting for response",
        ))
    }

    pub async fn notify(
        &self,
        method: &str,
        params: Option<serde_json::Value>,
    ) -> McpResult<()> {
        let notification = McpNotification::new(method, params);
        self.transport.send(McpMessage::Notification(notification)).await
    }
}

impl McpClient {
    pub async fn initialize(
        &self,
        implementation: McpImplementation,
        capabilities: McpClientCapabilities,
    ) -> McpResult<McpServerCapabilities> {
        let params = serde_json::json!({
            "clientInfo": implementation,
            "capabilities": capabilities,
            "protocolVersion": crate::protocol::LATEST_PROTOCOL_VERSION,
        });

        let response = self
            .request("initialize", Some(params))
            .await?;

        let server_capabilities: McpServerCapabilities = serde_json::from_value(response)?;
        *self.server_capabilites.write().await = Some(server_capabilities.clone());

        self.notify("initialize", None).await?;

        Ok(server_capabilities)
    }

    pub async fn tools_list(&self) -> McpResult<Vec<McpTool>> {
        let response = self
            .request("tools/list", None)
            .await?;
        // Get 'tools' in response, it's a Vec<McpTool>
        let tools = response.get("tools").ok_or(McpError::Other("Except tools!".into()))?;
        Ok(serde_json::from_value(tools.clone())?)
    }

    pub async fn tools_call(&self, tool_call: McpToolCall) -> McpResult<McpToolCallResult> {
        let params = serde_json::to_value(tool_call)?;
        let response = self
            .request("tools/call", Some(params))
            .await?;
        Ok(serde_json::from_value(response)?)
    }

    pub async fn shutdown(&self) -> McpResult<()> {
        self.request("shutdown", None).await?;
        self.notify("exit", None).await?;
        self.transport.close().await
    }
}

impl McpClient {
    pub async fn capabilities(&self) -> Option<McpServerCapabilities> {
        self.server_capabilites.read().await.clone()
    }
}